//==============================================================================
// HashSampleSearchEngine.cc
// ----------------------------------------------------------------------------
// A sample search engine based on hash 
//
//------------------------------------------------------------------------------
// $Id: $
//------------------------------------------------------------------------------
// $Log: $
//
//==============================================================================
//.......1.........2.........3.........4.........5.........6.........7.........8
//

#include "HashSampleSearchEngine.hh"

bool cmpBySourceLocation(const SampleSearchResult &a,const SampleSearchResult &b){
	return(a.mSourceLoc<b.mSourceLoc);
}

bool cmpBySourceIdx(const SampleSearchResult &a,const SampleSearchResult &b){
	return(a.mSourceIdx < b.mSourceIdx);
}


//-------------------------------------------
// * SortBySourceLoc
//-------------------------------------------
//
void 
SampleSearchResultVec::SortBySourceLoc(){
  stable_sort(this->begin(),this->end(),cmpBySourceLocation);
}

//-------------------------------------------
// * SortBySourceIdx
//-------------------------------------------
//
void 
SampleSearchResultVec::SortBySourceIdx(){
  stable_sort(this->begin(),this->end(),cmpBySourceIdx);
}


//-------------------------------------------
// * Create
//-------------------------------------------
// Creates a hash_map that can be used to quickly search for n-mers. 
//
void 
HashSampleSearchEngine::Create(Sequence &Seq){
  mTotalSamples = 0;
  mMap.resize(mBucketSize);
  SeqSampler SS;
  SequenceVector Fragments;

  long StartTime;
  

  //  SS.FragmentSeq(mSampleSize,(mSampleSize-1),Seq,Fragments);
  SS.FastFragmentSeq(mSampleSize,Seq,Fragments);
  // Insert all of the fragments.
  register int fsize = Fragments.size();
  for(int i = 0;i<fsize;i++){
	 if (mMap.Contains(Fragments[i])){
		int idx = mMap[Fragments[i]];
		if (mHashInfo[idx].mLocations.size() < mMaxLocsPerMer){
		  // i is the location of the ith fragment
		  mHashInfo[idx].AddHit(i,0);
		}else{
		  mbugMaxLocsCount++;
		}
	 }else{
		// Create a new hash info location
		mHashInfo.resize(mHashInfo.size()+1);
		int newIdx = mHashInfo.size()-1;
		mHashInfo[newIdx].AddHit(i,0);
		//mHashInfo[newIdx].mSequence = Fragments[i];
		// add hashinfo location to the map
		mMap[Fragments[i]] = newIdx;
	 }
  }
}

//Attempted speed up.  Has some bug. 

//-------------------------------------------
// * Add
//-------------------------------------------
// Adds the given sequence to a hash_map.  
// for n-mers.  
//
void  
HashSampleSearchEngine::Add(Sequence &Seq,int seqIdx){
  SeqSampler SS;
  mTotalSamples = 0;
  
  // Note, total bases will always be enough, and very conservative.   
  long newIdx = mHashInfo.size();
  mHashInfo.resize(mHashInfo.size()+Seq.size());  

  int fstart = 0;
  int fend = mSampleSize -1;
  do{
	 Sequence fragment;
	 Seq.GetSubSequence(fragment,fstart,fend);
	 
	 int idx;
	 if (mMap.Contains(fragment,idx)){                              
		if (mHashInfo[idx].mLocations.size() < mMaxLocsPerMer){
		  // fstart is the source index for the mer. 
		  mHashInfo[idx].AddHit(fstart,seqIdx);
		}else{
		  mbugMaxLocsCount++;
		}
	 }else{

		// Create a new hash info location
		mHashInfo[newIdx].AddHit(fstart,seqIdx);
		// add hashinfo location to the map
		mMap[fragment] = newIdx;
		newIdx++; // Point to the location where the next new mer will go. 
	 }

	 fstart++;
	 fend++;
	 
	 if (!(fend% 100000)){
		cout<<fend<<" bp in hash."<<endl;
	 }
	 
  }while(fend < Seq.size());
}  

/*

//-------------------------------------------
// * Add
//-------------------------------------------
// Adds the given sequence to a hash_map.  
// for n-mers.  
//
void  
HashSampleSearchEngine::Add(Sequence &Seq,int seqIdx){
  SeqSampler SS;
  mTotalSamples = 0;

  // Note, total bases will always be enough, and very conservative.   
  long newIdx = mHashInfo.size();
  mHashInfo.resize(mHashInfo.size()+Seq.size());  

  SequenceVector Fragments;
  //cout<<"Fast fragment seq..."<<flush;
  SS.FastFragmentSeq(mSampleSize,Seq,Fragments);
  //cout<<"done."<<endl;
  
  // Insert all of the fragments.  
  register int fsize = Fragments.size();
  for(int i = 0;i<fsize;i++){
	 int idx;
	 if (mMap.Contains(Fragments[i],idx)){                              
		if (mHashInfo[idx].mLocations.size() < mMaxLocsPerMer){
			 // i is the location of the ith fragment
			 mHashInfo[idx].AddHit(i,seqIdx);
		  }else{
			 mbugMaxLocsCount++;
		  }
	 }else{
		// Create a new hash info location
		//mHashInfo.resize(mHashInfo.size()+1);                           
		//int newIdx = mHashInfo.size()-1;
		mHashInfo[newIdx].AddHit(i,seqIdx);
		//mHashInfo[newIdx].mSequence = Fragments[i];
		// add hashinfo location to the map
		mMap[Fragments[i]] = newIdx;
		newIdx++; // Point to the location where the next new mer will go. 
	 }
  }
} 
*/

 
//-------------------------------------------
// * Create 
//-------------------------------------------
// Creates a hash_map for a set of sequences that can be quickly searched
// for n-mers.  
//
void  
HashSampleSearchEngine::Create(SequenceVector &SeqVec,int start, int end){
  //cout<<"mBucketSize="<<mBucketSize<<endl;
  mMap.resize(mBucketSize);
  SeqSampler SS;
  mTotalSamples = 0;

  // Pre-size as a guess
  long TotalBases = SeqVec.TotalBases();
  // Note, total bases will always be enough, and very conservative. 
  mHashInfo.resize(TotalBases);  

  long newIdx = 0;
  for(int sIdx = start;sIdx <=end;sIdx++){
	 SequenceVector *pFragments = new SequenceVector();
	 SS.FastFragmentSeq(mSampleSize,SeqVec[sIdx],*pFragments);            // 69
	 // Insert all of the fragments.  
	 register int fsize = (*pFragments).size();
	 for(int i = 0;i<fsize;i++){
		int idx;
		if (mMap.Contains((*pFragments)[i],idx)){                              // 72
		  //int idx = mMap[(*pFragments)[i]];		  
		  if (mHashInfo[idx].mLocations.size() < mMaxLocsPerMer){
			 // i is the location of the ith fragment
			 mHashInfo[idx].AddHit(i,sIdx);
		  }else{
			 mbugMaxLocsCount++;
		  }
		}else{
		  // Create a new hash info location
		  //mHashInfo.resize(mHashInfo.size()+1);                           
		  //int newIdx = mHashInfo.size()-1;
		  mHashInfo[newIdx].AddHit(i,sIdx);
		  //mHashInfo[newIdx].mSequence = Fragments[i];
		  // add hashinfo location to the map
		  mMap[(*pFragments)[i]] = newIdx;
		  newIdx++; // Point to the location where the next new mer will go. 
		}
	 }
	 delete pFragments;
  }
} 
 

// long start = (long)(sNum*sampleSpacing);
//	 long end = start+mSampleSize - 1;           // KJD WHAT??? Shouldn't this be mNumSamples? 

//---------------------------------------------------
// * SampleQuery
//---------------------------------------------------
// Written so that multiple calls tack samples onto end
// of SamplesOut
//
//
void 
HashSampleSearchEngine::SampleQuery(Sequence &QueryIn,SequenceVector &SamplesOut)
{
  long sampleSpacing = (long)(1000.0/(float) mSampleDensity);
  int NumSamples = (int)((float)QueryIn.size()/sampleSpacing);
  
  long sNumStart = SamplesOut.size();
  SamplesOut.resize(sNumStart+NumSamples);
  long qStart = 0;
  // For each sample
  for(long sNum = sNumStart;sNum < (sNumStart+NumSamples);sNum++){
	 // Make space for another sample
	 SamplesOut[sNum].resize(mSampleSize);
	 // copy the bases from the query to the sample out
	 for(int bNum = 0;bNum < mSampleSize;bNum++){
		SamplesOut[sNum][bNum]=QueryIn[bNum+qStart];
	 }
	 qStart+=sampleSpacing;
  }
}

//---------------------------------------------------
// * SlowSampleQuery
//---------------------------------------------------
//
//void 
//HashSampleSearchEngine::SlowSampleQuery(Sequence &QueryIn,SequenceVector &SamplesOut)
//{
// long sampleSpacing = (long)(1000.0/(float) mSampleDensity);
//  int NumSamples = (int)((float)QueryIn.size()/sampleSpacing);
  
//  for(long sNum = 0;sNum < NumSamples;sNum++){
//	 long start = (long)(sNum*sampleSpacing);
//	 long end = start+mSampleSize - 1;
//	 SamplesOut.AddSubSequence(QueryIn,start,end);
//  }
//}
long
HashSampleSearchEngine::Search(Sequence &Query,SampleSearchResultVec &Results){
  SequenceVector Queries;
  int bugSize = Query.size();
  Queries.push_back(Query);
  return(Search(Queries,Results));
}


long
HashSampleSearchEngine::SearchRev(Sequence &Query,SampleSearchResultVec &Results){
  SequenceVector Queries;
  int bugSize = Query.size();
  Queries.push_back(Query);
  return(SearchRev(Queries,Results));
}

//---------------------------------------------------
// * Search
//---------------------------------------------------
//
long
HashSampleSearchEngine::Search(SequenceVector &Queries,SampleSearchResultVec &Results){	 
  long sampleSpacing = (long)(1000.0/(float) mSampleDensity);
  ofstream bugout;
  long TotalSamplesThisSearch = 0;
  register int qsize = Queries.size();
  for(int q = 0;q< qsize;q++){
	 // Generate samples
	 SequenceVector QuerySamples;
	 // What about reverse compliments?? 
	 SampleQuery(Queries[q],QuerySamples);
	 register int qsampsize = QuerySamples.size();
	 TotalSamplesThisSearch+=qsampsize;
	 mTotalSamples+=qsampsize; 
	 
	 // Search each sample against source
	 for(int s = 0;s<qsampsize;s++){
		
		if (mMap.Contains(QuerySamples[s])){
		  int oldSize = Results.size();
		  long infoIdx = mMap[QuerySamples[s]];
		  int numHits = mHashInfo[infoIdx].mLocations.size();
		  // Create space for all the new results and then init them to proper values
		  Results.resize(oldSize+numHits);

		  for(int r = 0;r<numHits;r++){
			 int rIdx = oldSize+r;
			 Results[rIdx].mSourceIdx = mHashInfo[infoIdx].mSourceIdx[r];
			 Results[rIdx].mQueryIdx = q;
			 Results[rIdx].mSampleIdx = s;
			 Results[rIdx].mSourceLoc = mHashInfo[infoIdx].mLocations[r];
			 Results[rIdx].mSampleLoc = s*sampleSpacing;
		  }
		}
	 }
  }
  return(TotalSamplesThisSearch);
}




//---------------------------------------------------
// * SearchSingle
//---------------------------------------------------
// Searches the source for a single Query mer.  
// Query size much match the mer size for the search engine. 
//
void
HashSampleSearchEngine::SearchSingle(Sequence &Query,SampleSearchResultVec &Results){
  if (mMap.Contains(Query)){
	 int oldSize = Results.size();
	 long infoIdx = mMap[Query];
	 int numHits = mHashInfo[infoIdx].mLocations.size();

	 // Create space for all the new results and then init them to proper values
	 Results.resize(oldSize+numHits);
	 
	 for(int r = 0;r<numHits;r++){
		int rIdx = oldSize+r;
		Results[rIdx].mSourceIdx = mHashInfo[infoIdx].mSourceIdx[r];
		Results[rIdx].mQueryIdx = 0;
		Results[rIdx].mSampleIdx = 0;
		Results[rIdx].mSourceLoc = mHashInfo[infoIdx].mLocations[r];
		Results[rIdx].mSampleLoc = 0;
	 }
  }
}


//---------------------------------------------------
// * SearchRev
//---------------------------------------------------
//
long
HashSampleSearchEngine::SearchRev(SequenceVector &Queries,SampleSearchResultVec &Results){	 
  long sampleSpacing = (long)(1000.0/(float) mSampleDensity);
  ofstream bugout;

  long TotalSamplesThisSearch = 0;
  register int qsize = Queries.size();
  for(int q = 0;q< qsize;q++){
	 // Generate samples
	 SequenceVector QuerySamples;

	 Sequence RevQuery = Queries[q];
	 RevQuery.ReverseComplement();
	 SampleQuery(RevQuery,QuerySamples);

	 register int qsampsize = QuerySamples.size();
	 TotalSamplesThisSearch+=qsampsize;
	 mTotalSamples+=qsampsize; 
	 
	 // Search each sample against source
	 for(int s = 0;s<qsampsize;s++){
		
		if (mMap.Contains(QuerySamples[s])){
		  int oldSize = Results.size();
		  long infoIdx = mMap[QuerySamples[s]];
		  int numHits = mHashInfo[infoIdx].mLocations.size();
		  // Create space for all the new results and then init them to proper values
		  Results.resize(oldSize+numHits);

		  for(int r = 0;r<numHits;r++){
			 int rIdx = oldSize+r;
			 Results[rIdx].mSourceIdx = mHashInfo[infoIdx].mSourceIdx[r];
			 Results[rIdx].mQueryIdx = q;
			 Results[rIdx].mSampleIdx = s;
			 Results[rIdx].mSourceLoc = mHashInfo[infoIdx].mLocations[r];
			 Results[rIdx].mSampleLoc = s*sampleSpacing;
		  }
		}
	 }
  }
  return(TotalSamplesThisSearch);
}

//---------------------------------------------------
// * Search
//---------------------------------------------------
//
long
HashSampleSearchEngine::SearchWithSamples(SequenceVector &Samples,
														SampleSearchResultVec &Results){	 
  mTotalSamples = 0;	
  mTotalSamples+=Samples.size(); 
	 
  long sampleSpacing = (long)(1000.0/(float) mSampleDensity);

  // Search each sample against source
  register int ssize = Samples.size();
  for(int s = 0;s<ssize;s++){
		
	 if (mMap.Contains(Samples[s])){
		int oldSize = Results.size();
		long infoIdx = mMap[Samples[s]];
		int numHits = mHashInfo[infoIdx].mLocations.size();
		// Create space for all the new results and then init them to proper values
		Results.resize(oldSize+numHits);
		for(int r = 0;r<numHits;r++){
		  int rIdx = oldSize+r;
		  Results[rIdx].mSourceIdx = mHashInfo[infoIdx].mSourceIdx[r];
		  Results[rIdx].mQueryIdx = 0; 
		  Results[rIdx].mSampleIdx = s;
		  Results[rIdx].mSourceLoc = mHashInfo[infoIdx].mLocations[r];
		  Results[rIdx].mSampleLoc = s*sampleSpacing;
		}
	 }
  }
  return(Samples.size());
}
